#include "ASMInstruction.hpp"
#include "CodeGen.hpp"
#include "Module.hpp"

#include <iostream>
#include <memory>
#include <unordered_map>

void translate_main(CodeGen *codegen); // 将 main 函数翻译为汇编代码

int main() {
    auto *module = new Module();
    auto *codegen = new CodeGen(module);

    // 告诉汇编器将汇编放到代码段
    codegen->append_inst(".text", ASMInstruction::Atrribute);

    translate_main(codegen);

    std::cout << codegen->print();
    delete codegen;
    delete module;
    return 0;
}

// TODO: 按照提示补全
void translate_main(CodeGen *codegen) {
    std::unordered_map<std::string, int> offset_map;

    /* 声明 main 函数 */
    codegen->append_inst(".globl main", ASMInstruction::Atrribute);
    codegen->append_inst(".type main, @function", ASMInstruction::Atrribute);

    /* main 函数开始 */
    codegen->append_inst("main", ASMInstruction::Label);

    /* main 函数的 Prologue (序言) */
    codegen->append_inst("addi sp, sp, -96");
    codegen->append_inst("sd ra, 88(sp)");
    codegen->append_inst("sd s0, 80(sp)"); 
    codegen->append_inst("addi s0, sp, 96");

    /* main 函数的 label_entry */
    codegen->append_inst(".main_label_entry", ASMInstruction::Label);

    /* %op0 = alloca [10 x i32] */
    codegen->append_inst("%op0 = alloca [10 x i32]", ASMInstruction::Comment);
    // %op0 是一个指针，占8字节
    // 数组[10 x i32]占40字节(10 * 4)
    // 栈帧从s0向下分配空间
    offset_map["%op0"] = -16;  // 第一个变量在s0-16位置
    offset_map["array"] = -56; // 数组从s0-56开始，到s0-16结束
    codegen->append_inst("addi", {"t0", "s0", std::to_string(offset_map["array"])});
    codegen->append_inst("sd", {"t0", "s0", std::to_string(offset_map["%op0"])});

    /* %op1 = getelementptr [10 x i32], [10 x i32]* %op0, i32 0, i32 0 */
    codegen->append_inst(
        "%op1 = getelementptr [10 x i32], [10 x i32]* %op0, i32 0, i32 0",
        ASMInstruction::Comment);
    offset_map["%op1"] = -24; // 第二个变量在s0-24位置
    codegen->append_inst("ld", {"t0", "s0", std::to_string(offset_map["%op0"])});
    // 数组首地址就是%op1的值
    codegen->append_inst("sd", {"t0", "s0", std::to_string(offset_map["%op1"])});

    /* store i32 10, i32* %op1 */
    codegen->append_inst("store i32 10, i32* %op1", ASMInstruction::Comment);
    codegen->append_inst("ld", {"t0", "s0", std::to_string(offset_map["%op1"])});
    codegen->append_inst("li t1, 10");
    codegen->append_inst("sw t1, 0(t0)");

    /* %op2 = getelementptr [10 x i32], [10 x i32]* %op0, i32 0, i32 1 */
    codegen->append_inst(
        "%op2 = getelementptr [10 x i32], [10 x i32]* %op0, i32 0, i32 1",
        ASMInstruction::Comment);
    offset_map["%op2"] = -32; // 第三个变量在s0-32位置
    codegen->append_inst("ld", {"t0", "s0", std::to_string(offset_map["%op0"])});
    // 计算第二个元素的地址 = 数组首地址 + 4
    codegen->append_inst("addi t0, t0, 4");
    codegen->append_inst("sd t0, -32(s0)");

    /* %op3 = load i32, i32* %op1 */
    codegen->append_inst("%op3 = load i32, i32* %op1", ASMInstruction::Comment);
    offset_map["%op3"] = -36; // 第四个变量在s0-36位置(32位)
    codegen->append_inst("ld t0, -24(s0)");
    codegen->append_inst("lw t1, 0(t0)");
    codegen->append_inst("sw t1, -36(s0)");

    /* %op4 = mul i32 %op3, 3 */
    codegen->append_inst("%op4 = mul i32 %op3, 3", ASMInstruction::Comment);
    offset_map["%op4"] = -40; // 第五个变量在s0-40位置(32位)
    codegen->append_inst("lw t0, -36(s0)");
    codegen->append_inst("li t1, 3");
    codegen->append_inst("mul t2, t0, t1");
    codegen->append_inst("sw t2, -40(s0)");

    /* store i32 %op4, i32* %op2 */
    codegen->append_inst("store i32 %op4, i32* %op2", ASMInstruction::Comment);
    codegen->append_inst("lw t0, -40(s0)");
    codegen->append_inst("ld t1, -32(s0)");
    codegen->append_inst("sw t0, 0(t1)");

    /* %op5 = load i32, i32* %op2 */
    offset_map["%op5"] = -44; // 第六个变量在s0-44位置(32位)
    codegen->append_inst("ld t0, -32(s0)");
    codegen->append_inst("lw t1, 0(t0)");
    codegen->append_inst("sw t1, -44(s0)");

    /* ret i32 %op5 */
    codegen->append_inst("ret i32 %op5", ASMInstruction::Comment);
    codegen->append_inst("lw a0, -44(s0)");
    codegen->append_inst("j main_exit");

    /* main 函数的 Epilogue (收尾) */
    codegen->append_inst("main_exit", ASMInstruction::Label);
    codegen->append_inst("ld ra, 88(sp)");
    codegen->append_inst("ld s0, 80(sp)");
    codegen->append_inst("addi sp, sp, 96");
    codegen->append_inst("jr ra");
}